from __future__ import division
from pylab import *
from scipy import linalg
import sympy
from fJ_generator import fJ_generator
from scipy import optimize
def ooopti(fun, x0, args, Dfun):
    eps1 =     finfo(double).eps
    eps2 = finfo(double).eps
    tao = 0.000005
    
    k = 0
    v = 2
    x = asarray(x0)
    
    f = fun(x, args)
    J = Dfun(x, args)
    
    A = dot(J.T, J)
    g = dot(J.T, f)
    
    found = abs(g).max() < eps1
    u = tao * A.max()
    while (not found) and (k < 100):
        print k
        k += 1
        h = linalg.solve(A + u * eye(A.shape[0]), -g)
        if linalg.norm(h, 2) <= eps2 * (linalg.norm(x, 2) + eps2):
            found = True
            print 's'
        else:
            x_new = x + h
            ro = (fun(x, args) - fun(x + h, args)) / dot(0.5 * h.T, u * h - g)
            ro = ro.max()
            if ro > 0:
                x = x_new
                
                f = fun(x, args)
                J = Dfun(x, args)
    
                A = dot(J.T, J)
                g = dot(J.T, f)
                
                found = abs(g).max() < eps1
                if found == True:
                    print g
                u *= max(1 / 3, 1 - (2 * ro - 1) ** 3)
                v = 2
            else:
                u *= v
                v *= 2
    return x

if __name__ == '__main__':
    x0, x1, x2, x3, x4, x5 = sympy.symbols(['x0', 'x1', 'x2', 'x3', 'x4', 'x5'])
    t = sympy.Symbol('t')
    
    thecls = fJ_generator((x0 + x1 * t + x2 * t ** 2) * sympy.cos((x3 + x4 * t + x5 * t ** 2)), [x0, x1, x2, x3, x4, x5], {t:linspace(0, 1, 4410, endpoint=False)})
    data = thecls.f([20, -20, 1, 0, 20, 1])
    x = [3, 0, 3, 0, 25, -3]
    
#    (result, success, r1, r2, r3) = optimize.leastsq(thecls.f, x, args=(data), Dfun=thecls.J, full_output=True)
    result = ooopti(thecls.f, x, data, thecls.J)
    
    plot(thecls.f(result, zeros(4410)))
    plot(data, 'r--')
    plot(thecls.f(array(x), zeros(4410)), 'go')
    show()
